import cv2
import numpy as np
import pytesseract

lower_blue = np.array([100, 43, 46])
upper_blue = np.array([124, 255, 255])
lower_yellow = np.array([15, 55, 55])
upper_yellow = np.array([50, 255, 255])
lower_green = np.array([35, 43, 46])
upper_green = np.array([77, 255, 255])
lower = [lower_blue, lower_yellow, lower_green]
upper = [upper_blue, upper_yellow, upper_green]


# 读取图片文件
def imreadex(filename):
    return cv2.imdecode(np.fromfile(filename, dtype=np.uint8), cv2.IMREAD_COLOR)


# 避免出界
def point_limit(point):
    if point[0] < 0:
        point[0] = 0
    if point[1] < 0:
        point[1] = 0


# 预测车牌
def predict(car_pic):
    if type(car_pic) == type(""):
        raw_image = imreadex(car_pic)
    else:
        raw_image = car_pic

    # 调整大小
    raw_image = cv2.resize(raw_image, dsize=(1000, 750), interpolation=cv2.INTER_CUBIC)

    # 高斯滤波
    image = cv2.GaussianBlur(raw_image, (3, 3), 0)

    # 图片灰度化
    image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Sobel算子
    Sobel_x = cv2.Sobel(image, cv2.CV_16S, 1, 0)
    image = cv2.convertScaleAbs(Sobel_x)

    # 图像二值化
    ret, image = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU)

    # 闭操作
    kernelX = cv2.getStructuringElement(cv2.MORPH_RECT, (17, 5))
    image = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernelX)

    # 腐蚀膨胀
    kernelX = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 1))
    kernelY = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 16))

    image = cv2.dilate(image, kernelX)
    image = cv2.erode(image, kernelX)

    image = cv2.erode(image, kernelY)
    image = cv2.dilate(image, kernelY)

    # 中值滤波
    image = cv2.medianBlur(image, 15)

    # 查找轮廓
    contours = cv2.findContours(image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[1]

    car_contours = []
    for contour in contours:
        area = cv2.contourArea(contour)
        rect = cv2.boundingRect(contour)
        weight = rect[2]
        height = rect[3]
        if weight > (height * 2) and weight < (height * 5.5) and area > 2000:
            min_rect = cv2.minAreaRect(contour)
            car_contours.append(min_rect)

    card_imgs = []
    # 矩形区域可能是倾斜的矩形，需要仿射变换，不然OCR识别不了
    for rect in car_contours:
        if rect[2] > -1 and rect[2] < 1:  # 创造角度，使得左、高、右、低拿到正确的值
            angle = 1
        else:
            angle = rect[2]
        rect = (rect[0], (rect[1][0] + 5, rect[1][1] + 5), angle)  # 扩大范围，避免车牌边缘被排除

        box = cv2.boxPoints(rect)
        heigth_point = right_point = [0, 0]
        left_point = low_point = [1000, 750]
        for point in box:
            if left_point[0] > point[0]:
                left_point = point
            if low_point[1] > point[1]:
                low_point = point
            if heigth_point[1] < point[1]:
                heigth_point = point
            if right_point[0] < point[0]:
                right_point = point

        if left_point[1] <= right_point[1]:  # 正角度
            new_right_point = [right_point[0], heigth_point[1]]
            pts2 = np.float32([left_point, heigth_point, new_right_point])  # 字符只是高度需要改变
            pts1 = np.float32([left_point, heigth_point, right_point])
            M = cv2.getAffineTransform(pts1, pts2)
            dst = cv2.warpAffine(raw_image, M, (1000, 750))
            point_limit(new_right_point)
            point_limit(heigth_point)
            point_limit(left_point)
            if (int(left_point[1]) == int(heigth_point[1]) or int(left_point[0]) == int(new_right_point[0])):
                continue
            card_img = dst[int(left_point[1]):int(heigth_point[1]), int(left_point[0]):int(new_right_point[0])]
            card_imgs.append(card_img)
        elif left_point[1] > right_point[1]:  # 负角度

            new_left_point = [left_point[0], heigth_point[1]]
            pts2 = np.float32([new_left_point, heigth_point, right_point])  # 字符只是高度需要改变
            pts1 = np.float32([left_point, heigth_point, right_point])
            M = cv2.getAffineTransform(pts1, pts2)
            dst = cv2.warpAffine(raw_image, M, (1000, 750))
            point_limit(right_point)
            point_limit(heigth_point)
            point_limit(new_left_point)
            if (int(right_point[1]) == int(heigth_point[1]) or int(new_left_point[0]) == int(right_point[0])):
                continue
            card_img = dst[int(right_point[1]):int(heigth_point[1]), int(new_left_point[0]):int(right_point[0])]
            card_imgs.append(card_img)

    color = []
    if len(card_imgs) > 1:
        for i in range(3):
            for card_img in card_imgs:
                # 转换为HSV
                img_path = card_img
                hsv = cv2.cvtColor(img_path, cv2.COLOR_BGR2HSV)

                # 根据阈值构建掩膜
                mask = cv2.inRange(hsv, lower[i], upper[i])

                # 对原图像和掩膜进行位运算
                cv2.bitwise_and(img_path, img_path, mask=mask)
                count = 0
                height = mask.shape[0]
                width = mask.shape[1]
                for j in range(height):
                    for k in range(width):
                        if (mask[j][k] == 255):
                            count += 1

                if count / (height * width) > 0.5 or (count / (height * width) > 0.25 and i == 2):
                    color.append(card_img)
        card_imgs = color

    # 正常来说，只剩下一个，也就是车牌，其他都被过滤了
    card = np.zeros((350, 110, 3), np.uint8)
    text = ""
    if len(card_imgs) > 0:
        card = card_imgs[0]
        card = cv2.resize(card, dsize=(350, 110), interpolation=cv2.INTER_CUBIC)
        image = cv2.GaussianBlur(card, (3, 3), 0)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        ret, image = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU)
        image = image[10:, ]
        image = cv2.medianBlur(image, 3)
        kernel = np.ones((3, 3), np.uint8)
        image = cv2.erode(image, kernel, iterations=2)
        image = cv2.dilate(image, kernel, iterations=1)
        text = pytesseract.image_to_string(image, lang='chi_sim',config="--psm 7")
    text = text.replace(" ", "")
    text = text.replace(".", "")
    text = text.replace("-", "")
    text = text.replace("\n", "")
    if len(text) > 7 and (text[-1] == "1" or text[-1] == "|" or text[-1] == "I"):
        text = text[:-1]
    return card[10:, ], text


if __name__ == '__main__':
    print(predict("1.jpg"))
